'''
    关系：is_symptom（和is_accompany_with难区分）
        元素：
        A：疾病    B：症状
        句式：
        （A和B）/（B和A）有（啥/什么/没有/无）（关系/联系）？
        （得了/患上）A会（引起/造成）B嘛/吗/咩？
        （得了/患上）A是否会（引起/造成）B
        B是不是A的症状（之一）
        B是A的症状（之一）吗？

    关系：is_food
        元素：
        A：疾病    B：食物
        （得了/患上）A可以（吃/食用/饮食/摄入）B嘛/吗/咩？
        （得了/患上）A的病人可不可以（吃/食用/饮食/摄入）B
        （吃/食用/饮食/摄入）B对A（的（康复/缓解）/的病人）（有益/有好处/有不良影响/有坏处）嘛/吗/咩？
        B是A（患者/病人）的（忌食/宜吃）食物吗？

    关系：is_belonging
        元素：
        A：疾病    B：科室
        B（可以/能够/能不能/可不可以）治疗A吗
        （得了/患上）A（可以到/一般是去）B（去）（治疗/挂号/看病)吗
        我可以去B挂号看A吗
        如果我得了A，可以到B挂号看病咩
        A和B啥关系？

    关系：is_recommend_drug
        元素：
        A：疾病    B：药品
        B是不是治A的药物？
        B可以治疗A吗？
        B能不能治A啊？
        A的推荐药物是B吗？
        B是A的推荐药品嘛？
        得了A推荐吃B吗？
        A是吃B吗？
        一般A推荐吃B嘛

    关系：is_need_check
        元素：
        A：疾病    B：检查
        B可以诊断出A吗？
        B能够（查出/测出/试出）A嘛
        B是A的检查项目吗？
        B可不可以检查出A吗？
        A要做B吗？



'''
import os
import random
from nlpcda.tools.Simbert import Simbert

# 单位样本扩充数量（n表示一个样本扩充n个相似样本）
num_of_expand = 1


class generate_data:
    def gen_sentence(self, label, As, Bs, tempaltes, simbert, num=1, is_enhancement=True):
        '''
        :param label: 标签
        :param As: 实体A列表
        :param Bs: 实体B列表
        :param templates: 句子模板列表
        :param num：要生成的句子数量
        :return: list like [{label:a,text:b}...]
        '''
        result = []
        for i in range(num):
            A = random.choice(As)
            tempalte = random.choice(tempaltes)

            if Bs != []:
                B = random.choice(Bs)
                text = tempalte.format(A, B)
            else:
                text = tempalte.format(A)

            if is_enhancement:
                synonyms = simbert.replace(sent=text, create_num=num_of_expand)
                for item in synonyms:
                    result.append({'label': label, 'text': item[0]})

            result.append({'label': label, 'text': text})

        return result

    def save_data(self, list_for_dict, save_path):
        fileObject = open(save_path, 'w')
        fileObject.write('label\ttext\n')
        for dictionary in list_for_dict:
            fileObject.write(dictionary['label'] + '\t' + dictionary['text'] + '\n')
        fileObject.close()


if __name__ == '__main__':
    cur_dir = '/'.join(os.path.abspath(__file__).split('/')[:-1])
    # 特征词路径
    disease_path = os.path.join(cur_dir, 'dict/disease.txt')
    symptom_path = os.path.join(cur_dir, 'dict/symptom.txt')
    food_path = os.path.join(cur_dir, 'dict/food.txt')
    department_path = os.path.join(cur_dir, 'dict/department.txt')
    check_path = os.path.join(cur_dir, 'dict/check.txt')
    drug_path = os.path.join(cur_dir, 'dict/drug.txt')
    # 加载特征词
    disease_wds = [i.strip() for i in open(disease_path, encoding='UTF-8') if i.strip()]
    symptom_wds = [i.strip() for i in open(symptom_path, encoding='UTF-8') if i.strip()]
    food_wds = [i.strip() for i in open(food_path, encoding='UTF-8') if i.strip()]
    department_wds = [i.strip() for i in open(department_path, encoding='UTF-8') if i.strip()]
    check_wds = [i.strip() for i in open(check_path, encoding='UTF-8') if i.strip()]
    drug_wds = [i.strip() for i in open(drug_path, encoding='UTF-8') if i.strip()]

    # 设置各类别句子模板
    is_symptom_templates = [
        '{0}和{1}有什么关系？', '{0}和{1}有啥关系啊', '{0}和{1}有没有关系？', '{0}和{1}有什么联系',
        '{1}和{0}有什么关系', '{1}和{0}有啥关系？', '{1}和{0}有没有关系？', '{1}和{0}有什么联系',
        '得了{0}会有{1}？', '得了{0}一般会引起{1}这种症状嘛？', '患上{0}的人会有{1}的症状咩',
        '{0}会不会造成{1}？', '得了{0}会有{1}的症状嘛？', '一个人如果得了{0}会有{1}吗？',
        '{1}是不是{0}的症状之一', '{1}是不是{0}的症状', '{1}是{0}的症状之一吗'
    ]

    is_food_templates = [
        '得了{0}可以吃{1}吗？', '得了{0}可以食用{1}吗？', '患上{0}可以饮用{1}嘛', '得了{0}能够摄入{1}吗？', '患了{0}可以食用{1}咩？',
        '得了{0}的病人可以摄入{1}吗？', '患上{0}的病人可不可以吃{1}', '{1}对{0}的康复或缓解有影响吗？', '食用{1}对得了{0}的病人有好处吗',
        '摄入{1}对{0}的康复有益吗', '食用{1}对{0}的康复有害吗', '{1}对{0}的缓解有不良影响嘛？', '{1}是{0}患者的忌食食物吗？',
        '{1}是{0}患者的宜吃食物吗？', '{1}是{0}病人的宜吃食品吗', '{1}适合患了{0}的人吃吗？'
    ]

    is_belonging_templates = [
        '{1}可以治疗{0}吗？', '{1}能够治疗{0}嘛', '{1}能不能治疗{0}', '{1}可不可以治{0}', '得了{0}可以到{1}去治疗吗？',
        '得了{0}可以去{1}挂号看病吗？', '患上{0}一般是到{1}看医生吗', '我可以去{1}挂号看{0}吗', '如果我得了{0}，是不是到{1}去挂号看病？',
        '假如一个人得了{0}，可以到{1}看吗', '{0}和{1}是啥关系', '{0}和{1}有什么联系', '{1}和{0}有什么关系'
    ]

    is_recommend_drug_templates = [
        '{1}是不是治{0}的药物', '{1}可以治疗{0}吗？', '{1}能不能治{0}？', '{0}的推荐药品是{1}嘛', '{1}是{0}的推荐药吗？',
        '得了{0}推荐吃{1}吗', '得了{0}可以吃{1}咩？', '{0}是吃{1}吗', '一般{0}推荐吃{1}吗', '如果我得了{0}，是不是应该吃{1}治疗',
        '{1}对{0}有治疗效果嘛？'
    ]

    is_need_check_templates = [
        '{1}可以诊断出{0}吗', '{1}能够查出{0}吗？', '{1}能够试出{0}吗', '一般做{1}是否能查出{0}？', '{1}可以检测出{0}嘛？',
        '{1}是{0}的检查项目吗', '{1}可不可以检查出{0}', '我得了{0}，可以通过{1}检测出来吗', '得了{0}要做{1}吗'
    ]

    others_templates = [
        '{0}有啥症状？', '{0}有什么症状？', '{0}的症状是什么？', '什么是{0}', '{0}是什么？', '{0}有什么并发症吗？',
        '{0}的并发症是啥', '{0}的病因一般是啥？', '{0}的宜吃食物是什么？', '得了{0}建议吃点啥？', '{0}的忌食食品有啥？', '得了{0}的人不能吃什么？', '{0}病人最好不要吃什么？',
        '吃什么食品对患了{0}的病人有好处', '得了{0}一般要吃啥药？', '{0}患者一般要做什么检查', '{0}做什么检查项目能够查出来？', '{0}有什么预防措施',
        '{0}的治愈率是多少？', '{0}一般会引起什么症状？', '{0}的目标人群？', '{0}有什么推荐药物吗', '得了{0}一般到哪个科室挂号？', '得了{0}一般到哪里看？',
        '{0}的易感人群是哪些人？'
    ]

    others_templates_food = [
        '{0}对治疗什么病有好处', '{0}是啥病的推荐食物？', '得了什么病适合吃{0}', '{0}对什么病的康复有阻碍', '啥病推荐吃{0}'
    ]

    other_templates_check = [
        '{0}一般能查出什么病？', '得了什么病可以通过{0}测出来', '{0}可以检查出啥病症？'
    ]

    other_templates_department = [
        '{0}一般治疗那些疾病？', '一般到{0}看啥病？', '什么病人是去{0}挂号？'
    ]

    other_templates_drug = [
        '{0}是啥病的推荐药物？', '{0}是什么病的推荐药品', '{0}可以治疗什么病？', '得了什么病推荐吃{0}'
    ]

    other_templates_symptom = [
        '{0}的症状一般在什么病当中会出现？', '得了什么病会有{0}的症状', '{0}是啥病的症状？', '什么病会出现{0}'
    ]

    generator = generate_data()
    # train
    print('start...train set!')
    config = {
        'model_path': '/home/cike/diploma project/code_reference/PaddleHub/demo/my_relation_classification/utils/nlpcda/tools/models/chinese_simbert_L-12_H-768_A-12',
        'CUDA_VISIBLE_DEVICES': '6',
        'max_len': 64,
        'seed': 1
    }
    simbert = Simbert(config=config)
    result_is_symptom = generator.gen_sentence('is_symptom', disease_wds, symptom_wds, is_symptom_templates, simbert,
                                               100, True)
    result_is_food = generator.gen_sentence('is_food', disease_wds, food_wds, is_food_templates, simbert, 100, True)
    result_is_belonging = generator.gen_sentence('is_belonging', disease_wds, department_wds, is_belonging_templates,
                                                 simbert, 100, True)
    result_is_recommend_drug = generator.gen_sentence('is_need_check', disease_wds, check_wds, is_need_check_templates,
                                                      simbert, 100, True)
    result_is_need_check = generator.gen_sentence('is_recommend_drug', disease_wds, drug_wds,
                                                  is_recommend_drug_templates,
                                                  simbert, 100, True)
    result_other = generator.gen_sentence('other', disease_wds, [], others_templates, simbert,
                                          len(others_templates) * 8,
                                          True) + generator.gen_sentence('other', check_wds, [], other_templates_check,
                                                                         simbert, len(other_templates_check) * 8,
                                                                         True) + generator.gen_sentence(
        'other', department_wds, [], other_templates_department, simbert, len(other_templates_department) * 8,
        True) + generator.gen_sentence(
        'other', drug_wds, [], other_templates_drug, simbert, len(other_templates_drug) * 8,
        True) + generator.gen_sentence(
        'other', symptom_wds, [], other_templates_symptom, simbert, len(other_templates_symptom) * 8,
        True) + generator.gen_sentence(
        'other', food_wds, [], others_templates_food, simbert, len(others_templates_food) * 8, True)
    result_all = result_is_symptom + result_is_food + result_other
    random.shuffle(result_all)
    print('train set length:{0}'.format(len(result_all)))
    generator.save_data(result_all, '../data/enhanced/train.txt')

    # eval
    config = {
        'model_path': '/home/cike/diploma project/code_reference/PaddleHub/demo/my_relation_classification/utils/nlpcda/tools/models/chinese_simbert_L-12_H-768_A-12',
        'CUDA_VISIBLE_DEVICES': '6',
        'max_len': 64,
        'seed': 2
    }
    simbert = Simbert(config=config)
    result_is_symptom = generator.gen_sentence('is_symptom', disease_wds, symptom_wds, is_symptom_templates, simbert,
                                               30, True)
    result_is_food = generator.gen_sentence('is_food', disease_wds, food_wds, is_food_templates, simbert, 30, True)
    result_is_belonging = generator.gen_sentence('is_belonging', disease_wds, department_wds, is_belonging_templates,
                                                 simbert, 30, True)
    result_is_recommend_drug = generator.gen_sentence('is_need_check', disease_wds, check_wds, is_need_check_templates,
                                                      simbert, 30, True)
    result_is_need_check = generator.gen_sentence('is_recommend_drug', disease_wds, drug_wds,
                                                  is_recommend_drug_templates,
                                                  simbert, 30, True)
    result_other = generator.gen_sentence('other', disease_wds, [], others_templates, simbert,
                                          len(others_templates) * 2,
                                          True) + generator.gen_sentence('other', check_wds, [], other_templates_check,
                                                                         simbert, len(other_templates_check) * 2,
                                                                         True) + generator.gen_sentence(
        'other', department_wds, [], other_templates_department, simbert, len(other_templates_department) * 2,
        True) + generator.gen_sentence(
        'other', drug_wds, [], other_templates_drug, simbert, len(other_templates_drug) * 2,
        True) + generator.gen_sentence(
        'other', symptom_wds, [], other_templates_symptom, simbert, len(other_templates_symptom) * 2,
        True) + generator.gen_sentence(
        'other', food_wds, [], others_templates_food, simbert, len(others_templates_food) * 2, True)
    result_all = result_is_symptom + result_is_food + result_other
    random.shuffle(result_all)
    print('dev set length:{0}'.format(len(result_all)))
    generator.save_data(result_all, '../data/enhanced/dev.txt')

    # test
    config = {
        'model_path': '/home/cike/diploma project/code_reference/PaddleHub/demo/my_relation_classification/utils/nlpcda/tools/models/chinese_simbert_L-12_H-768_A-12',
        'CUDA_VISIBLE_DEVICES': '6',
        'max_len': 64,
        'seed': 3
    }
    simbert = Simbert(config=config)
    result_is_symptom = generator.gen_sentence('is_symptom', disease_wds, symptom_wds, is_symptom_templates, simbert,
                                               10, True)
    result_is_food = generator.gen_sentence('is_food', disease_wds, food_wds, is_food_templates, simbert, 10, True)
    result_is_belonging = generator.gen_sentence('is_belonging', disease_wds, department_wds, is_belonging_templates,
                                                 simbert, 10, True)
    result_is_recommend_drug = generator.gen_sentence('is_need_check', disease_wds, check_wds, is_need_check_templates,
                                                      simbert, 10, True)
    result_is_need_check = generator.gen_sentence('is_recommend_drug', disease_wds, drug_wds,
                                                  is_recommend_drug_templates,
                                                  simbert, 10, True)
    result_other = generator.gen_sentence('other', disease_wds, [], others_templates, simbert,
                                          len(others_templates) * 2,
                                          True) + generator.gen_sentence('other', check_wds, [], other_templates_check,
                                                                         simbert, len(other_templates_check) * 2,
                                                                         True) + generator.gen_sentence(
        'other', department_wds, [], other_templates_department, simbert, len(other_templates_department) * 2,
        True) + generator.gen_sentence(
        'other', drug_wds, [], other_templates_drug, simbert, len(other_templates_drug) * 2,
        True) + generator.gen_sentence(
        'other', symptom_wds, [], other_templates_symptom, simbert, len(other_templates_symptom) * 2,
        True) + generator.gen_sentence(
        'other', food_wds, [], others_templates_food, simbert, len(others_templates_food) * 2, True)
    result_all = result_is_symptom + result_is_food + result_other
    random.shuffle(result_all)
    print('test set length:{0}'.format(len(result_all)))
    generator.save_data(result_all, '../data/enhanced/test.txt')

    # for item in result_is_symptom:
    #     print(item)
    #
    # print('-----------------------------------------')
    #
    # for item in result_is_food:
    #     print(item)
    #
    # print('-----------------------------------------')
    #
    # for item in result_other:
    #     print(item)
